
#include "common.h"
#include "page.h"
#include "kheap.h"

#include "scrn.h"

//end is defined in the linker script.
extern u32int end;
extern page_directory_t* kerneldir;


u32int placement_addr =(u32int)&end;
static heap_t* kernelheap = 0;

#define MAGIC_NUM  0xDEADDEAD
#define HEAP_HEAD_SIZE (sizeof(heap_head_t))
#define HEAP_FOOT_SIZE (sizeof(heap_foot_t))
#define HEAP_TAG_SIZE (HEAP_HEAD_SIZE+HEAP_FOOT_SIZE)
#define HEAP_BLK_MIN_SIZE 1

#define GET_FOOT(head) (heap_foot_t*)((u32int)(head)+(head)->size+HEAP_HEAD_SIZE)

#define CHECK_BLK(head) { \
          heap_foot_t* pf = GET_FOOT(head);\
          ASSERT(pf->magic == MAGIC_NUM); \
          ASSERT(pf->phead == (u32int)(head));\
        }

#define DUMP_BLK(head) {\
         puts("\nHEAD:");puts_hex((u32int)(head));\
         puts(" size:");puts_dec(((heap_head_t*)(head))->size);\
         puts(" used:");puts_dec(((heap_head_t*)(head))->used);\
         puts("\n");}

static heap_head_t* find_unused_blk(heap_t* head,u32int sz,u32int align);

static heap_head_t* mergblk(heap_t* pheap,heap_head_t* ph)
{
   ASSERT(ph);
   
   heap_foot_t* pf = GET_FOOT(ph);
   ASSERT(pf->magic == MAGIC_NUM);
   
   //merg right
   if((u32int)pf+HEAP_FOOT_SIZE < pheap->end)
   { 
        heap_head_t* rph = (heap_head_t*)((u32int)pf+HEAP_FOOT_SIZE);
        ASSERT(rph->magic == MAGIC_NUM);
        if(!rph->used)
        {
           heap_foot_t* rpf = GET_FOOT(rph);
           ASSERT(rpf->magic == MAGIC_NUM);
           ph->size += rph->size+HEAP_TAG_SIZE;
           rpf->phead =(u32int)ph;
           pf = rpf;
           CHECK_BLK(ph);
        }
   } 
   //merg left
   if((u32int)ph != pheap->start)
   {
       heap_foot_t* lpf=(heap_foot_t*)((u32int)ph-HEAP_FOOT_SIZE);
       ASSERT(lpf->magic == MAGIC_NUM);
       heap_head_t* lph =(heap_head_t*)lpf->phead;
       ASSERT(lph->magic == MAGIC_NUM);
       CHECK_BLK(lph);

       if(!lph->used)
       {
          lph->size += ph->size+ HEAP_TAG_SIZE;
          pf->phead =(u32int)lph;
          ph = lph;
          CHECK_BLK(lph);
    //      puts("merg left\n");
       }
   }
   CHECK_BLK(ph);

   return ph;
}

void kfree(u32int p)
{
  ASSERT(kernelheap);
  ASSERT(p> kernelheap->start);
  ASSERT(p< kernelheap->end);

  heap_head_t* ph = (heap_head_t*)(p-HEAP_HEAD_SIZE);
  ASSERT(ph->magic == MAGIC_NUM);

  heap_foot_t* pf = (heap_foot_t*)(p+ph->size);
  ASSERT(pf->magic == MAGIC_NUM);
  ASSERT(pf->phead == (u32int)ph);

  ph = mergblk(kernelheap,ph);
  ph->used =0;  
}
/*
 * expanded size by page
 * */
int expandheap(heap_t* heap,u32int sz,page_directory_t* dir)
{
   ASSERT((heap->end%0x1000)==0);
   u32int i=0,idx=0;
   
   if(heap->end+sz > heap->start+heap->maxlen)
       return 0;
   
   if(sz%0x1000 == 0)
     idx = sz/0x1000;
      else
        idx = sz/0x1000+1;

   while(i<idx)
   {
      //assume all the page for heap has been allocated before page
      //enable.
      alloc_kernel_frame(get_page(heap->end+i*0x1000,0,dir));
      i++;
   }  
   heap->end += idx*0x1000;
   return idx*0x1000;
}
/*
 * find a unused blk which(not include TAG) is larger than sz
 * */
static heap_head_t* find_unused_blk(heap_t* head,u32int sz,u32int align)
{
    ASSERT(head);
    heap_head_t* ph = head->head;
   
    while(ph)
    {
        ASSERT(ph->magic == MAGIC_NUM);

        if((!ph->used)&&(ph->size>=sz))
        {
            if(!align) 
                return ph;
             else
             {
                if((u32int)ph/0x1000 < ((u32int)ph+ph->size) /0x1000)
                {    
                    s32int more = (0x1000-(u32int)ph)%0x1000 - HEAP_HEAD_SIZE;
                    if(0 == more)
                    {//aligned HEAP.
                        return ph;
                    }else
                    {   //if it can be split
                        if((more>(HEAP_TAG_SIZE+HEAP_BLK_MIN_SIZE))
                         &&(ph->size-more > sz))
                        {//big enough after split;
                            return ph;
                        }
                     }   
                }
        
             }
         }


        if(((u32int)ph+HEAP_TAG_SIZE)>=head->end)
            break;

        u32int next = (u32int)ph+ph->size+HEAP_TAG_SIZE;
        if(next >= head->end)
            next = 0;
        ph=(heap_head_t*)next;
    }
    return 0;    
}

void init_heap(u32int start,u32int end,u32int maxlen)
{
    kernelheap = (heap_t*)kmalloc(sizeof(heap_t));

    kernelheap->head = 0;
    kernelheap->start = start;
    kernelheap->end = start;    
    kernelheap->maxlen = maxlen;
    u32int sz = expandheap(kernelheap,end-start,kerneldir);
    ASSERT(sz == end - start);

    kernelheap->head = (heap_head_t*)start;
    kernelheap->head->magic = MAGIC_NUM;
    kernelheap->head->used = 0;
    kernelheap->head->size = end - start-HEAP_TAG_SIZE;
    kernelheap->end = end;
    heap_foot_t* pf = (heap_foot_t*)(end-HEAP_FOOT_SIZE);
    pf->magic = MAGIC_NUM;
    pf->phead = (u32int)kernelheap->head;
    CHECK_BLK(kernelheap->head);
}

/*
 * if align = 1, return the aligned one.
 * */
heap_head_t* splitblk(heap_head_t* ph,u32int sz,u32int align)
{
   ASSERT(ph);
   ASSERT(ph->magic == MAGIC_NUM);
   ASSERT(!ph->used);

// align
   if(align)
   {
      ASSERT((u32int)ph/0x1000 < ((u32int)ph+ph->size)/0x1000);
      s32int more = (0x1000-(u32int)ph)%0x1000 - HEAP_HEAD_SIZE;
      if(more <0)
          more += 0x1000;

      if(more > 0)//need split
      {
        puts("more:");
        puts_hex(more);
        ASSERT(more >(HEAP_TAG_SIZE+HEAP_BLK_MIN_SIZE));
        u32int oldsz = ph->size;
        ph->size = more - HEAP_TAG_SIZE;
        heap_foot_t* npf = GET_FOOT(ph);
        npf->magic = MAGIC_NUM;
        npf->phead = (u32int)ph;
        CHECK_BLK(ph);
        ph = (heap_head_t*)((u32int)npf+HEAP_FOOT_SIZE);
        ph->magic = MAGIC_NUM;
        ph->size = oldsz - more;
        npf = GET_FOOT(ph);
        ASSERT(npf->magic == MAGIC_NUM);
        npf->phead =(u32int)ph;
        CHECK_BLK(ph);       
      }
   }

//make new after the returned blk.
   u32int left = ph->size - sz - HEAP_TAG_SIZE;
   if(left > HEAP_BLK_MIN_SIZE)
   {
     u32int oldsz = ph->size;
     heap_head_t* nph = 0;
     heap_foot_t* oldpf = (heap_foot_t*)(((u32int)ph)+HEAP_HEAD_SIZE+oldsz);
     ASSERT(oldpf->magic == MAGIC_NUM);
     ASSERT(oldpf->phead ==(u32int)ph);

     heap_foot_t* pf = (heap_foot_t*)(((u32int)ph)+HEAP_HEAD_SIZE+sz);
     pf->magic = MAGIC_NUM;
     pf->phead = (u32int)ph;
     ph->size = sz;
     nph =  (heap_head_t*)((u32int)pf+HEAP_FOOT_SIZE);
     nph->magic = MAGIC_NUM;
     nph->size = left;
     nph->used = 0;
     oldpf->phead = (u32int)nph;
     //sanity test
     ASSERT((u32int)oldpf == (((u32int)nph)+nph->size+HEAP_HEAD_SIZE));
     CHECK_BLK(nph);
     CHECK_BLK(ph);
   } 
   return ph;
}

static u32int alloc(heap_t* pheap,
        u32int sz, 
        u32int align,
        page_directory_t*  dir)
{
   ASSERT(pheap);
   
   if (sz == 0) return 0;

   heap_head_t* ph = find_unused_blk(pheap,sz,align);
   if(!ph)
   {
       u32int end = pheap->end;
       u32int pagealloc;
       if(align)//if align, alloc 1 more page to split later.
           pagealloc = expandheap(pheap,sz+HEAP_TAG_SIZE+0x1000,dir);
       else 
            pagealloc = expandheap(pheap,sz+HEAP_TAG_SIZE,dir);

       ph = (heap_head_t*)end;
       ph->size = pagealloc-HEAP_TAG_SIZE;
       ph->magic = MAGIC_NUM;
       ph->used = 0;
       heap_foot_t* pf = (heap_foot_t*)(pheap->end-HEAP_FOOT_SIZE);
       pf->magic = MAGIC_NUM;
       pf->phead = (u32int)ph;
   } 
   
   ph = splitblk(ph,sz,align);
   ph->used = 1;
   
   CHECK_BLK(ph);
   return (u32int)ph + HEAP_HEAD_SIZE;
}

u32int kmalloc_int(u32int sz, int align,u32int* phys)
{
   if(kernelheap)
   {
       return alloc(kernelheap,sz,align,kerneldir);
   }else
   {
        if (align&&(placement_addr&0xFFFFF000 ))
        {
            if(placement_addr%0x1000)
            {
                placement_addr &= 0xFFFFF000;
                placement_addr += 0x1000;
            }
        }
        if (phys)
        {
            *phys = placement_addr;   
        }
        u32int tmp = placement_addr;
        placement_addr += sz;
        return tmp;
   }
}

u32int kmalloc_a(u32int sz)
{
  return kmalloc_int(sz,1,0);
}

u32int kmalloc_p(u32int sz, u32int *phys)
{
  return kmalloc_int(sz,0,phys);
}

u32int kmalloc_ap(u32int sz, u32int *phys)
{
  return kmalloc_int(sz,1,phys);
}

u32int kmalloc(u32int sz)
{
  return kmalloc_int(sz,0,0);
}


